import json
import os
import re

import xlrd
import click
import numpy as np
import pandas as pd

from . import TYPE_DEBT, TYPE_BENEFIT, TYPE_CASH
from .download import download_csv, download_bonus_payment
from .utils import read_process_config


class DataProcess:
    def __init__(self, code, read_dir, files):
        self.read_dir = read_dir
        self.files = files
        self.config = read_process_config()
        self.dataset = {}
        self.code = code

    def read_data(self, key):
        """从i问财下载的企业资产负债表、利润表以及现金流量表读取财报数据
           key: pre_process.json 中定义的键，包含 debt、benefit、cash
        """
        filename = self.files[key]
        workbook = xlrd.open_workbook(os.path.join(self.read_dir, filename), ignore_workbook_corruption=True)
        df = pd.read_excel(workbook, sheet_name='Worksheet', header=1, index_col=0)
        # 删除空行
        df.dropna(inplace=True)
        # 将省略符号 '--' 替换为 0
        df.replace({'--': 0.0}, inplace=True)
        # 取最近 6 年的数据，正序排列
        df = df.T[:6][::-1]
        df.columns.name = ''

        config = read_process_config()
        miss_auto = []
        for field in config[key]['auto']:
            if field not in df:
                miss_auto.append(field)
                df[field] = np.zeros(6)

        if miss_auto:
            miss_str = '\n'.join(miss_auto)
            click.echo(f"{miss_str}\n{'-' * 60}")
            click.secho('以上字段未找到，自动赋值为 0 ！\n', fg='yellow')

        hand_fields = config[key].get('hands')
        if hand_fields:
            hand_str = '\n'.join(hand_fields)
            click.echo(f"{hand_str}\n{'-' * 60}")
            click.secho(f'以上字段需要手动赋值，请稍后在 correct.csv 中录入！\n', fg='red')

        correct_fields = config[key].get('correct')
        if correct_fields:
            correct_str = '\n'.join(correct_fields)
            click.echo(f"{correct_str}\n{'-' * 60}")
            click.secho(f'以上字段或许需要调整，请稍后在 correct.csv 中修改！\n', fg='red')

        self.dataset[key] = df
        return df

    def save_data(self):
        # 水平方向拼接多个数据集
        df_all = pd.concat(self.dataset.values(), axis=1)
        df_all.to_csv(os.path.join(self.read_dir, f'{self.code}.csv'))
        return df_all


def init_project(code, download, input_dir):
    from .pre_process import DataProcess

    if download:
        for file_type in (TYPE_DEBT, TYPE_BENEFIT, TYPE_CASH):
            download_csv(file_type, code, input_dir)
        download_bonus_payment(code, download_dir=input_dir)

    files = {}
    file_pattern = re.compile(code + r'_(\w+)_year.xls')
    for file in os.listdir(input_dir):
        if result := file_pattern.match(file):
            files[result.group(1)] = file

    click.echo(json.dumps(files, indent=2))

    if len(files.keys()) < 3:
        click.secho(f'三大报表不齐全：资产负债表、利润表及现金流量表！', fg='red')
        click.echo(f'请使用 --download 选项自动完成数据的下载或将下载好的报表手动放入 {input_dir} 目录！')
        return

    click.echo(f'有关 {code} 的文件准备就绪，即将开始初始化！\n')

    process = DataProcess(code, input_dir, files)
    years = {}
    sheet, year = None, None
    for file_type in files:
        df = process.read_data(file_type)
        years[file_type] = df.index.tolist()

        if year is None and sheet is None:
            sheet, year = file_type, years[file_type]
            correct_df = pd.DataFrame(index=year)
        elif year != years[file_type]:
            click.secho(f'财务报表 {files[file_type]} 和 {files[sheet]} 的年份信息不一致!', fg='red')
            click.echo(f'{files[file_type]} : {years[file_type]!r}')
            click.echo(f'{files[sheet]} : {year!r}')
            return

        for correct_field in process.config[file_type].get('correct', []):
            correct_df[correct_field] = df[correct_field]

        for hand_field in process.config[file_type].get('hands', []):
            correct_df[hand_field] = [np.NAN for _ in correct_df.index]
            if file_type == TYPE_DEBT:
                correct_df.loc[correct_df.index > 2018, ['其他流动资产里的理财产品', '其他流动资产里的结构性存款']] = 0.0
                if correct_df[correct_df.index <= 2018].index.tolist():
                    correct_df.loc[correct_df.index <= 2018, '应收款项融资'] = 0.0

    process.save_data()
    bonus_path = os.path.join(input_dir, 'bonus.csv')
    if os.path.exists(bonus_path):
        bonus = pd.read_csv(bonus_path, index_col='报告期')
        correct_df = pd.merge(correct_df, bonus, how='left', left_index=True, right_index=True)
        correct_df['分红总额'].fillna(0.0, inplace=True)
    else:
        correct_df['分红总额'] = [np.NAN for _ in year]

    correct_path = os.path.join(input_dir, 'correct.csv')
    if os.path.exists(correct_path):
        if click.prompt(f'input/{code}/correct.csv 已经存在，是否覆盖？(y/n)', 'y').lower() == 'n':
            return
    correct_df.to_csv(correct_path)
    click.secho(f'\n输出 {correct_path} 文件成功！', fg='green', bg='white')


def output_csv_file(code, input_csv, output_csv):
    """generate merged CSV in output dir"""
    data = pd.read_csv(input_csv, index_col=0)
    correct = pd.read_csv(os.path.join('input', code, 'correct.csv'), index_col=0, thousands=',')
    correct.fillna(0.0, inplace=True)

    data.drop(columns=correct.columns, inplace=True, errors='ignore')
    data.merge(correct, left_index=True, right_index=True).to_csv(output_csv)
    click.secho(f'\n输出 {output_csv} 文件成功！', fg='green', bg='white')

